#!/bin/env python
import os
from netCDF4 import Dataset
import numpy as np
from mpl_toolkits.basemap import Basemap, interp
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib.colors


#--Directory Setting
input_besttrack_dir = "./input_data/besttrack" # Best track TCG data
input_gpi_dir = "./output_data" # GPI data

output_dir = "./pic_data" # Output for figures


undef = -9.99E33 # undefined value

#--Input
#--DGPI
f1 = Dataset("%s/dgpi_climmonthly.nc" % (input_gpi_dir), "r")
dgpi = f1.variables['dgpi'][:]

#--ENGPI
f2 = Dataset("%s/engpi_climmonthly.nc" % (input_gpi_dir), "r")
engpi = f2.variables['engpi'][:]

#--Observed TCG
f3 = Dataset("%s/tcg_climmonthly.nc" % (input_besttrack_dir), "r")
tcg = f3.variables['tcg'][:]

#--Longitude Latitude
lon_in = f1.variables['lon']
lat_in = f1.variables['lat']

lons,lats = np.meshgrid(lon_in[:],lat_in[:])

tmax,jmax,imax = np.shape(dgpi)

print (np.shape(dgpi), np.shape(engpi), np.shape(tcg))

#--Normalized 
for tt in range(tmax):
    dgpi[tt,:,:]  = tcg[tt,:,:].sum()/dgpi[tt,:,:].sum() * dgpi[tt,:,:]
    engpi[tt,:,:] = tcg[tt,:,:].sum()/engpi[tt,:,:].sum() * engpi[tt,:,:]


#-Plot
fig=plt.figure(num=None, figsize=(12,8), dpi=150, facecolor='w', edgecolor='k')

#--Color map
cmap = cm.jet

#--Contours
contours=np.array([0.1,0.2,0.4,0.6,0.8,1,1.5,2,4,6,8,10,15,20])
norm = matplotlib.colors.BoundaryNorm(contours,cmap.N)

#--plot TCG
ax=plt.subplot(3,1,1)

m = Basemap(projection='merc',llcrnrlat=-40,urcrnrlat=45, llcrnrlon=30,urcrnrlon=360, lat_ts=5, resolution='c')
m.drawcoastlines()
m.drawmeridians(np.arange(-30,360,30),labels=[0,0,0,1],fontsize=6)
m.drawparallels(np.arange(-40,40,10),labels=[1,0,0.0],fontsize=6)

x,y=m(lons,lats)
m.contourf(x,y,tcg.sum(axis=0),contours, cmap=cmap, norm=norm)
m.colorbar(drawedges=False)
plt.title('(a) TCG (Merged Reanalsyis, 1980-2017)', fontsize=16)

#--plot ENGPI
ax=plt.subplot(3,1,2)

m = Basemap(projection='merc',llcrnrlat=-40,urcrnrlat=45, llcrnrlon=30,urcrnrlon=360, lat_ts=5, resolution='c')
m.drawcoastlines()
m.drawmeridians(np.arange(-30,360,30),labels=[0,0,0,1],fontsize=6)
m.drawparallels(np.arange(-40,40,10),labels=[1,0,0.0],fontsize=6)

x,y=m(lons,lats)
m.contourf(x,y,engpi.sum(axis=0),contours, cmap=cmap, norm=norm)
m.colorbar(drawedges=False)
plt.title('(b) ENGPI (Merged Reanalysis, 1980-2017)', fontsize=16)

#--plot DGPI
ax=plt.subplot(3,1,3)

m = Basemap(projection='merc',llcrnrlat=-40,urcrnrlat=45, llcrnrlon=30,urcrnrlon=360, lat_ts=5, resolution='c')
m.drawcoastlines()
m.drawmeridians(np.arange(-30,360,30),labels=[0,0,0,1],fontsize=6)
m.drawparallels(np.arange(-40,40,10),labels=[1,0,0.0],fontsize=6)

x,y=m(lons,lats)
m.contourf(x,y,dgpi.sum(axis=0),contours, cmap=cmap, norm=norm)
m.colorbar(drawedges=False)
plt.title('(c) DGPI (Merged Reanalysis, 1980-2017)', fontsize=16)


if not os.path.exists(output_dir):
    os.makedirs(output_dir)

outpng="%(outdir)s/Fig.png" %{"outdir":output_dir}
plt.tight_layout(pad=1.0, w_pad=1.0, h_pad=1.0)

plt.savefig(outpng, dpi=100, bbox_inches='tight')
